package com.spx.chapter09;

import com.spx.chapter05.pojo.Event;
import com.spx.util.ClickSource;
import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;

/**
 * create by undeRdoG on  2022-05-06  23:03
 * 凡心所向，素履以往，生如逆旅，一苇以航。
 */
public class BufferingSinkExample {

    public static void main(String[] args) {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        SingleOutputStreamOperator<Event> stream = env.addSource(new ClickSource())
                .assignTimestampsAndWatermarks(
                        WatermarkStrategy.<Event>forBoundedOutOfOrderness(Duration.ZERO)
                                .withTimestampAssigner(new SerializableTimestampAssigner<Event>() {
                                    @Override
                                    public long extractTimestamp(Event element, long recordTimestamp) {
                                        return element.timestamp;
                                    }
                                })
                );


        // 批量缓存输出
        stream.addSink(new BufferingSink(10L));
    }

    public static class BufferingSink implements SinkFunction<Event>, CheckpointedFunction {

        private Long bufferingSize;

        private List<Event> bufferedElements;

        public BufferingSink(Long bufferingSize) {
            this.bufferingSize = bufferingSize;
            this.bufferedElements = new ArrayList<>();
        }


        // 定义算子状态
        private ListState<Event> checkpointedState;


        @Override
        public void invoke(Event value, Context context) throws Exception {
            this.bufferedElements.add(value);
            if (bufferedElements.size() == bufferingSize){
                // 模拟进行批量写入
                for (Event bufferedElement : bufferedElements) {
                    System.out.println(bufferedElement);
                }
                bufferedElements.clear();
            }
        }

        // 持久化时调用的方法
        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {

            checkpointedState.clear();

            // 对状态进行持久化，即把 bufferedElements 中的数据 放到 checkpointedState 中
            for (Event bufferedElement : bufferedElements) {
                checkpointedState.add(bufferedElement);
            }
        }

        // 初始化或者故障恢复时调用
        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {

            // 定义算子状态
            ListStateDescriptor<Event> bufferedStateDescriptor = new ListStateDescriptor<>("bufferedState", Event.class);

            // 根据算子状态的操作句柄拿到算子状态
            this.checkpointedState = context.getOperatorStateStore().getListState(bufferedStateDescriptor);

            // 从故障的情况下恢复,需要将ListState中的所有元素复制到本地变量bufferedElements来恢复状态
            if (context.isRestored()){
                for (Event element : this.checkpointedState.get()){
                    this.bufferedElements.add(element);
                }
            }
        }
    }
}
